package org.gaaidou.ptarmigan.springboot.data.jpa.datasource;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.*;

public class RoutingDataSource extends AbstractRoutingDataSource {

    private final Random random = new Random();

    private final Map<String, List<String>> multipleRouteKeyMap = new HashMap<>();


    public void multipleKey(String key) {
        if (key.matches(".*-[0-9]")) {
            String keyPrefix = key.substring(0, key.lastIndexOf("-"));
            if (!multipleRouteKeyMap.containsKey(keyPrefix)) {
                multipleRouteKeyMap.put(keyPrefix, new LinkedList<>());
            }
            multipleRouteKeyMap.get(keyPrefix).add(key);
        }

    }

    @Override
    protected Object determineCurrentLookupKey() {
        String current = RouteStrategyHolder.current();
        if (null != current) {
            if (multipleRouteKeyMap.containsKey(current)) {
                int i = random.nextInt(multipleRouteKeyMap.get(current).size());
                return multipleRouteKeyMap.get(current).get(i);
            } else {
                return current;
            }
        }
        return null;
    }
}

